package org.testing.dbunit.extension.xml;

import static com.google.common.base.Preconditions.checkNotNull;

import java.io.ByteArrayInputStream;
import java.nio.charset.Charset;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;

import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.xml.FlatXmlDataSet;
import org.dbunit.dataset.xml.FlatXmlDataSetBuilder;
import org.xml.sax.InputSource;

public class DataSetBuilder {
    private List<Table> schemas = new ArrayList<Table>();

    public void add(Schema schema) {
	Collection<Table> tables = schema.getTables();
	for (Table table : tables) {
	    schemas.add(checkNotNull(table));
	}
    }

    @Deprecated
    // FIXME should add just complete schemas
    public void add(Table table) {
	schemas.add(checkNotNull(table));
    }

    @Override
    public String toString() {
	StringBuilder dataset = new StringBuilder(
		"<!DOCTYPE dataset SYSTEM \"testing.dtd\">");
	dataset.append("<dataset>");
	for (Table schema : this.schemas) {
	    dataset.append(schema.toString());
	}
	dataset.append("</dataset>");
	return dataset.toString();
    }

    public byte[] getBytes() {
	return build().getBytes(Charset.forName("UTF-8"));
    }

    public void empty(Class<? extends Table> emptyTable) {
	schemas.add(new EmptyTable(emptyTable.getSimpleName()));
    }

    public IDataSet flat() {
	try {
	    FlatXmlDataSet build = new FlatXmlDataSetBuilder()//
		    .setDtdMetadata(false)//
		    .build(new InputSource(//
			    new ByteArrayInputStream(this.getBytes())));
	    return build;
	} catch (Exception e) {
	    throw new RuntimeException(e);
	}
    }

    public String build() {
	return toString();
    }
}